import pika
import threading
import time
import logging
import json  # For potential future use if message body becomes complex

# --- 配置 ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(threadName)s - %(message)s')
logger = logging.getLogger(__name__)

RABBITMQ_HOST = 'localhost'
RABBITMQ_PORT = 15672
RABBITMQ_USER = 'guest'
RABBITMQ_PASS = '123456'
MAIN_EXCHANGE_NAME = 'tasks_main_exchange'  # 与发布者一致
EXCHANGE_TYPE = 'topic'  # 与发布者一致

# 定义我们要监听的任务类型及其对应的队列和路由键绑定
# 路由键应与 FastAPI 中 TaskConfig 的 mq_routing_key_prefix 匹配
# 使用 topic exchange 的通配符
TASK_QUEUES_AND_ROUTING_KEYS = {
    # 队列名: 绑定键 (路由键模式)
    "q_erp_tpl_sales_orders": "erp.tpl.sales_orders",  # 精确匹配
    "q_tpl_wms_shipments": "tpl.wms.shipments",  # 精确匹配
    "q_wms_tpl_inventory": "wms.tpl.inventory",  # 精确匹配
    "q_other_tasks": "other.#",  # 匹配所有 other 开头的
}


# 模拟的任务执行逻辑 (实际应用中会更复杂)
def execute_erp_to_tpl_sales_order_sync(task_config_id: str):
    logger.info(f"[ERP->TPL] 开始同步销售订单, 任务ID: {task_config_id}")
    # 1. 获取此 task_config_id 的详细配置 (DB连接, SQL等) - 可能从数据库或本地配置
    # 2. 连接ERP数据库, 执行查询
    # 3. 连接TPL数据库, 执行插入/更新
    time.sleep(3)  # 模拟工作
    logger.info(f"[ERP->TPL] 销售订单同步完成, 任务ID: {task_config_id}")


def execute_tpl_to_wms_shipment_sync(task_config_id: str):
    logger.info(f"[TPL->WMS] 开始同步发货单, 任务ID: {task_config_id}")
    time.sleep(4)  # 模拟工作
    logger.info(f"[TPL->WMS] 发货单同步完成, 任务ID: {task_config_id}")


def execute_wms_to_tpl_inventory_sync(task_config_id: str):
    logger.info(f"[WMS->TPL] 开始同步库存, 任务ID: {task_config_id}")
    time.sleep(2)  # 模拟工作
    logger.info(f"[WMS->TPL] 库存同步完成, 任务ID: {task_config_id}")


def execute_other_task(task_config_id: str, routing_key: str):
    logger.info(f"[OTHER] 开始执行其他任务, 任务ID: {task_config_id}, 来源路由: {routing_key}")
    time.sleep(1)  # 模拟工作
    logger.info(f"[OTHER] 其他任务执行完成, 任务ID: {task_config_id}")


# 将路由键前缀映射到处理函数
TASK_HANDLERS = {
    "erp.tpl.sales_orders": execute_erp_to_tpl_sales_order_sync,
    "tpl.wms.shipments": execute_tpl_to_wms_shipment_sync,
    "wms.tpl.inventory": execute_wms_to_tpl_inventory_sync,
    # "other.#" 这样的通配符不能直接作为字典的key, 我们需要在回调中处理
}


def message_callback(ch, method, properties, body):
    task_id = ""
    try:
        task_id = body.decode('utf-8')
        routing_key = method.routing_key  # 获取消息实际的路由键
        logger.info(f"收到任务. 队列: '{method.exchange}'(实际上是队列名)->'{routing_key}', 任务ID: '{task_id}'")

        # 根据路由键选择处理函数
        handler = TASK_HANDLERS.get(routing_key)  # 精确匹配

        if handler:
            handler(task_id)
        elif routing_key.startswith("other."):  # 处理通配符匹配到的
            execute_other_task(task_id, routing_key)
        else:
            logger.warning(f"未找到路由键 '{routing_key}' 对应的处理函数。任务ID: '{task_id}'")
            # 也许应该 nack 并发送到死信队列

        ch.basic_ack(delivery_tag=method.delivery_tag)
        logger.info(f"任务 '{task_id}' (路由: {routing_key}) 已处理并 ACK.")

    except Exception as e:
        logger.error(f"处理任务 '{task_id}' (路由: {method.routing_key if method else 'N/A'}) 失败: {e}", exc_info=True)
        # 考虑 NACK 逻辑
        # ch.basic_nack(delivery_tag=method.delivery_tag, requeue=False) # 发送到DLX或丢弃


def consume_from_queue(queue_name: str, binding_key: str):
    thread_name = threading.current_thread().name
    logger.info(f"线程 '{thread_name}' 启动: 消费队列 '{queue_name}' (绑定键: '{binding_key}')")

    while True:  # 添加循环以支持重连
        try:
            credentials = pika.PlainCredentials(RABBITMQ_USER, RABBITMQ_PASS)
            parameters = pika.ConnectionParameters(RABBITMQ_HOST, RABBITMQ_PORT, '/', credentials)
            # parameters = pika.ConnectionParameters(host=RABBITMQ_HOST, port=RABBITMQ_PORT, heartbeat=60)  # 添加心跳
            connection = pika.BlockingConnection(parameters)
            channel = connection.channel()

            channel.exchange_declare(exchange=MAIN_EXCHANGE_NAME, exchange_type=EXCHANGE_TYPE, durable=True)

            # 声明队列 (durable=True)
            channel.queue_declare(queue=queue_name, durable=True)

            # 绑定队列
            channel.queue_bind(exchange=MAIN_EXCHANGE_NAME, queue=queue_name, routing_key=binding_key)
            logger.info(
                f"线程 '{thread_name}': 队列 '{queue_name}' 已绑定到 '{MAIN_EXCHANGE_NAME}' (key: '{binding_key}')")

            channel.basic_qos(prefetch_count=1)  # 每次处理一条消息
            channel.basic_consume(queue=queue_name, on_message_callback=message_callback, auto_ack=False)

            logger.info(f"线程 '{thread_name}': [*] 等待来自队列 '{queue_name}' 的消息...")
            channel.start_consuming()

        except pika.exceptions.AMQPConnectionError as e:
            logger.error(f"线程 '{thread_name}': RabbitMQ 连接错误: {e}. 10秒后重试...")
            time.sleep(10)
        except pika.exceptions.ChannelClosedByBroker:
            logger.error(f"线程 '{thread_name}': Channel被Broker关闭. 10秒后重试...")
            time.sleep(10)
        except pika.exceptions.ChannelWrongStateError:
            logger.error(f"线程 '{thread_name}': Channel状态错误. 10秒后重试...")
            time.sleep(10)
        except KeyboardInterrupt:
            logger.info(f"线程 '{thread_name}': 消费者停止 (KeyboardInterrupt).")
            if 'connection' in locals() and connection.is_open:
                connection.close()
            break  # 退出重试循环
        except Exception as e:
            logger.error(f"线程 '{thread_name}': 未捕获的异常: {e}", exc_info=True)
            if 'connection' in locals() and connection.is_open:
                connection.close()
            time.sleep(10)  # 发生其他未知错误也重试
            # break # 或者直接退出


if __name__ == "__main__":
    consumer_threads = []
    logger.info("启动 RabbitMQ 消费者...")

    for queue, routing_key_pattern in TASK_QUEUES_AND_ROUTING_KEYS.items():
        thread = threading.Thread(
            target=consume_from_queue,
            args=(queue, routing_key_pattern),
            name=f"Consumer-{queue}"
        )
        consumer_threads.append(thread)
        thread.daemon = True  # 设置为守护线程，主线程退出时它们也会退出
        thread.start()
        logger.info(f"为队列 '{queue}' (路由: '{routing_key_pattern}') 启动消费线程: {thread.name}")

    try:
        while True:  # 保持主线程运行，除非有 KeyboardInterrupt
            time.sleep(1)
    except KeyboardInterrupt:
        logger.info("主线程收到中断信号，正在关闭...")
    finally:
        logger.info("消费者程序退出。")